﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LPLib
{
    /// <summary>
    /// Abstract class for Markov models of nucleotide evolution.
    /// </summary>
    public abstract class NucleotideEvolutionModel
    {
        public NucleotideEvolutionModel(double time, double[] parameters = null)
        {
            this.time = time;
            if (parameters != null)
            {
                this.parameters = parameters;
            }
            else
            {
                this.parameters = new double[0];
            }
        }

        /// <summary>
        /// The probability that a nucleotide inititially in state originNucleotide has state destinationNucleotide at the specified time.
        /// </summary>
        /// <param name="originNucleotide">The original nucleotide state.</param>
        /// <param name="destinationNucleotide">The state of the nucleotide at the observation time.</param>
        /// <returns>The probability that originNucleotide is destinationNucleotide at the (default) observation time.</returns>
        public abstract double TransitionProbability(int originNucleotide, int destinationNucleotide);
        /// <summary>
        /// The probability that a nucleotide inititially in state originNucleotide has state destinationNucleotide at the specified time.
        /// </summary>
        /// <param name="originNucleotide">The original nucleotide state.</param>
        /// <param name="destinationNucleotide">The state of the nucleotide at the observation time.</param>
        /// <param name="_time">The branch length over which the observation occurs.</param>
        /// <returns>The probability that originNucleotide is destinationNucleotide at the observation time.</returns> 
        public abstract double TransitionProbability(int originNucleotide, int destinationNucleotide, double _time);
        
        public abstract double dTransitionProbabilityDT(int originNucleotide, int destinationNucleotide, double _time);

        public abstract double[] TypeProbabilities(double _distance);

        public abstract double[] DestinationPMF(double[] originPMF, double _time);

        public abstract double[] DestinationPMF(double[] originPMF);

        protected double time;
        protected double[] parameters;
        public abstract double Time { get; set; }
        public double[] Parameters { get { return parameters; } set { } }

    }

    /// <summary>
    /// Implements the Jukes-Cantor evolution model.
    /// </summary>
    public class JCEvolutionModel : NucleotideEvolutionModel
    {
        private double onDiagonal;
        private double offDiagonal;
        private double p;
        public override double Time
        {
            get { return time; }
            set
            {
                time = value;
                p = 0.75 * (1 - Math.Exp(-time));
                onDiagonal = 1 - p;
                offDiagonal = p / 3;
            }
        }

        public JCEvolutionModel(double _time)
            : base(_time)
        {
            time = _time;
            p = 0.75 * (1 - Math.Exp(-_time));
            onDiagonal = 1 - p;
            offDiagonal = p / 3;
            parameters = new double[1];
            parameters[0] = 1;
        }

        public override double TransitionProbability(int originNucleotide, int destinationNucleotide)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            return originNucleotide == destinationNucleotide ? onDiagonal : offDiagonal;
        }

        public override double TransitionProbability(int originNucleotide, int destinationNucleotide, double _time)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            double _p = 0.75 * (1 - Math.Exp(-_time));
            return originNucleotide == destinationNucleotide ? 1 - _p : _p / 3;
        }

        public override double dTransitionProbabilityDT(int originNucleotide, int destinationNucleotide, double _time)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            if (originNucleotide == destinationNucleotide)
            {
                return -0.75 * Math.Exp(-_time);
            }
            else
            {
                return 0.25 * Math.Exp(-_time);
            }
        }

        public static double ProbabilityToTime(double probability)
        {
            return -Math.Log(1 - 1.333333 * probability);
        }

        public override double[] TypeProbabilities(double _time)
        {
            double _p = 0.75 * (1 - Math.Exp(- _time));
            return new double[] {1-_p, _p, 2 * _p };
        }

        /// <summary>
        /// Computes the PMF resulting from Juke-Cantor evolution over the default time acting on a given initial state.
        /// </summary>
        /// <param name="originNucleotide">The initial PMF.</param>
        /// <returns>The PMF resulting from evolution over time from the initial PMF.</returns>
        public override double[] DestinationPMF(double[] originPMF)
        {
            // TODO: test to see if originNucleotide is a gap

            double[] destinationPMF = new double[originPMF.Length];
            if (originPMF.Length > 4)
            {
                destinationPMF[4] = originPMF[4];
            }
            for (int i = 0; i < 4; i++)
            {
                for (int j = 0; j < 4; j++)
                {
                    if (i == j) // on-diagonal
                    {
                        destinationPMF[i] += onDiagonal * originPMF[j];
                    }
                    else
                    {
                        destinationPMF[i] += offDiagonal * originPMF[j];
                    }
                }
            }
            return destinationPMF;
        }

        /// <summary>
        /// Computes the PMF resulting from Jukes-Cantor evolution over time acting on a given initial state.
        /// </summary>
        /// <param name="originNucleotide">The initial PMF.</param>
        /// <param name="_time">The evolutionary time.</param>
        /// <returns>The PMF resulting from evolution over time _time from the initial PMF.</returns>        
        public override double[] DestinationPMF(double[] originPMF, double _time)
        {
            // TODO: test to see if originNucleotide is a gap

            double _p = 0.75 * (1 - Math.Exp(-_time));
            double _onDiagonal = 1 - _p;
            double _offDiagonal = _p / 3;

            double[] destinationPMF = new double[originPMF.Length];
            if (originPMF.Length > 4)
            {
                destinationPMF[4] = originPMF[4];
            }

            for (int i = 0; i < 4; i++)
            {
                for (int j = 0; j < 4; j++)
                {
                    if (i == j)
                    {
                        destinationPMF[i] += _onDiagonal * originPMF[j];
                    }
                    else
                    {
                        destinationPMF[i] += _offDiagonal * originPMF[j];
                    }
                }
            }
            return destinationPMF;
        }

    }

    /// <summary>
    /// Implements the Kimura evolution model.
    /// </summary>
    public class KimuraEvolutionModel : NucleotideEvolutionModel
    {
        private double onDiagonal;
        private double offDiagonalTransition;
        private double offDiagonalTransversion;
        private double tau;
        public double Kappa;
        public override double Time
        {
            get { return time; }
            set
            {
                time = value;
                tau = 2 * time / (Kappa + 1);
                offDiagonalTransversion = 0.25 * (1 - Math.Exp(-tau));
                offDiagonalTransition = 0.25 * (1 + Math.Exp(-tau) - 2 * Math.Exp(-time));
                onDiagonal = 1 - offDiagonalTransition - 2 * offDiagonalTransversion;
            }
        }

        public KimuraEvolutionModel(double _time, double[] kappa)
            : base(_time, kappa)
        {
            Kappa = kappa[0];
            time = _time;
            tau = 2 * time / (Kappa + 1);
            offDiagonalTransversion = 0.25 * (1 - Math.Exp(-tau));
            offDiagonalTransition = 0.25 * (1 + Math.Exp(-tau) - 2 * Math.Exp(-time));
            onDiagonal = 1 - offDiagonalTransition - 2 * offDiagonalTransversion;
        }

        /// <summary>
        /// Computes the PMF resulting from Kimura evolution over the defaultBranchLength acting on a given initial pure state.
        /// </summary>
        /// <param name="originNucleotide">The initial PMF.</param>
        /// <param name="mu">The evolutionary distance mu.</param>
        /// <returns>The PMF resulting from evolution over distance mu from the initial PMF.</returns>
        public override double[] DestinationPMF(double[] originPMF)
        {
            // TODO: test to see if originNucleotide is a gap

            double[] destinationPMF = new double[originPMF.Length];
            if (originPMF.Length > 4)
            {
                destinationPMF[4] = originPMF[4];
            }
            for (int i = 0; i < 4; i++)
            {
                for (int j = 0; j < 4; j++)
                {
                    if (i == j) // on-diagonal
                    {
                        destinationPMF[i] += onDiagonal * originPMF[j];
                    }
                    else if (i / 2 == j / 2) // transition
                    {
                        destinationPMF[i] += offDiagonalTransition * originPMF[j];
                    }
                    else
                    {
                        destinationPMF[i] += offDiagonalTransversion * originPMF[j];
                    }
                }
            }
            return destinationPMF;
        }

        /// <summary>
        /// Computes the PMF resulting from Kimura evolution over distance branchLength acting on a given initial pure state.
        /// </summary>
        /// <param name="originNucleotide">The initial PMF.</param>
        /// <param name="mu">The evolutionary distance mu.</param>
        /// <returns>The PMF resulting from evolution over distance mu from the initial PMF.</returns>        
        public override double[] DestinationPMF(double[] originPMF, double _time)
        {
            // TODO: test to see if originNucleotide is a gap

            double _tau = 2 * _time / (Kappa + 1);
            double x = Math.Exp(-_tau);
            double _offDiagonalTransversion = (1 - x)/4;
            double _offDiagonalTransition = (1 + x - 2 * Math.Exp(-_time)) / 4;
            double _onDiagonal = 1 - _offDiagonalTransition - 2 * _offDiagonalTransversion;

            double[] destinationPMF = new double[originPMF.Length];
            if (originPMF.Length > 4)
            {
                destinationPMF[4] = originPMF[4];
            }

            for (int i = 0; i < 4; i++)
            {
                for (int j = 0; j < 4; j++)
                {
                    if (i == j)
                    {
                        destinationPMF[i] += _onDiagonal * originPMF[j];
                    }
                    else if (i / 2 == j / 2)
                    {
                        destinationPMF[i] += _offDiagonalTransition * originPMF[j];
                    }
                    else
                    {
                        destinationPMF[i] += _offDiagonalTransversion * originPMF[j];
                    }
                }
            }
            return destinationPMF;
        }

        public override double TransitionProbability(int originNucleotide, int destinationNucleotide)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            if (originNucleotide == destinationNucleotide)
            {
                return onDiagonal;
            }
            else if (originNucleotide / 2 == destinationNucleotide / 2)
            {
                return offDiagonalTransition;
            }
            else
            {
                return offDiagonalTransversion;
            }
        }

        public override double TransitionProbability(int originNucleotide, int destinationNucleotide, double _time)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            double x = Math.Exp(-2 * _time / (Kappa + 1));
            if (originNucleotide == destinationNucleotide) // identity
            {
                return (1 + x + 2 * Math.Exp(-_time)) / 4;
            }
            else if (originNucleotide / 2 == destinationNucleotide / 2) // transition
            {
                return (1 + x - 2 * Math.Exp(-_time)) / 4;
            }
            else // transversion
            {
                return (1 - x) / 4;
            }
        }

        //public double TransitionProbability(byte originNucleotide, byte destinationNucleotide)
        //{
        //    // return zero if either is a gap
        //    if (originNucleotide > 3 || destinationNucleotide > 3)
        //        return 0;

        //    if (originNucleotide == destinationNucleotide)
        //    {
        //        return onDiagonal;
        //    }
        //    else if (originNucleotide / 2 == destinationNucleotide / 2)
        //    {
        //        return offDiagonalTransition;
        //    }
        //    else
        //    {
        //        return offDiagonalTransversion;
        //    }
        //}

        //public double TransitionProbability(byte originNucleotide, byte destinationNucleotide, double _time)
        //{
        //    // return zero if either is a gap
        //    if (originNucleotide > 3 || destinationNucleotide > 3)
        //        return 0;

        //    double x = Math.Exp(-2 * _time / (Kappa + 1));
        //    if (originNucleotide == destinationNucleotide)
        //    {
        //        return (1 + x - 2 * Math.Exp(-_time)) / 4;
        //    }
        //    else if (originNucleotide / 2 == destinationNucleotide / 2)
        //    {
        //        return (1 + x - 2 * Math.Exp(-_time)) / 4;
        //    }
        //    else
        //    {
        //        return (1 - x) / 4;
        //    }
        //}

        public override double dTransitionProbabilityDT(int originNucleotide, int destinationNucleotide, double _time)
        {
            // return zero if either is a gap
            if (originNucleotide > 3 || destinationNucleotide > 3)
                return 0;

            double x = Math.Exp(-2 * _time / (Kappa + 1));
            double y = Math.Exp(-_time);
            if (originNucleotide == destinationNucleotide) // identity
            {
                return -(y + x / (Kappa + 1)) / 2; 
            }
            else if (originNucleotide / 2 == destinationNucleotide / 2) // transition
            {
                return (y - x / (Kappa + 1) ) / 2;
            }
            else // transversion
            {
                return x / (2*(Kappa + 1));
            }
        }

        public override double[] TypeProbabilities(double _time)
        {
            double _tau = 2 * _time / (Kappa + 1);
            double x = Math.Exp(-_tau);
            double _TwoTimesOffDiagonalTransversion = 0.5 * (1 - x);
            double _offDiagonalTransition = 0.25 * (1 + x - 2 * Math.Exp(-_time));
            double _onDiagonal = 1 - _offDiagonalTransition - _TwoTimesOffDiagonalTransversion;
            return new double[] { _onDiagonal, _offDiagonalTransition, _TwoTimesOffDiagonalTransversion};
        }
    }
}
